from sklearn.metrics import confusion_matrix, roc_curve, auc, classification_report
from sklearn.preprocessing import label_binarize
import pandas as pd
import matplotlib.pyplot as plt
from .utils import Inspector, write_dict_to_file, zip_utils
import torch
import pathlib
import numpy as np
from typing import List


class Result:
    def __init__(self, **kwargs):
        self.kwargs = kwargs

    def put(self, key, value):
        self.kwargs[key] = value

    def update(self, dic):
        self.kwargs.update(dic)

    def get(self, key):
        assert key in self.kwargs, f"key {key} not found in result"
        return self.kwargs[key]


class Processor:
    def __init__(self):
        self.inspector = Inspector(self)
        self.inspector.inspect(self.message)
        self.inspector.inspect(self.post_message)
        self.inspector.inspect(self.last_message)

    def message(self):
        # input settings
        # output g_data
        pass

    def post_message(self):
        # input settings + result + g_data
        # output None
        pass

    def last_message(self):
        # input settings + result + g_data
        # output None
        pass

    def distribute_message(self, coll_dict: dict):
        kwargs = self.inspector.distribute("message", coll_dict)
        out = self.message(**kwargs)
        return out if isinstance(out, dict) else {}

    def distribute_post_message(self, coll_dict: dict):
        kwargs = self.inspector.distribute("post_message", coll_dict)
        self.post_message(**kwargs)

    def distribute_last_message(self, coll_dict: dict):
        kwargs = self.inspector.distribute("last_message", coll_dict)
        self.last_message(**kwargs)


class ArrayProcessor(Processor):
    def __init__(self):
        super().__init__()
        self.array = []

    def append(self, value):
        self.array.append(value)


class ClassificationReportMetrics(Processor):
    def __init__(self):
        super().__init__()
        self.best_acc = None
        self.acc_array, self.precision_array, self.recall_array, self.f1_array = [], [], [], []

    def message(self, y_true, y_pred, idx_to_class, metrics_average, classification_report_kwargs) -> dict:
        target_names = [idx_to_class[i] for i in range(len(idx_to_class))]
        report_dict = classification_report(
            y_true, y_pred, target_names=target_names, output_dict=True, **classification_report_kwargs)
        cur_acc, precision, recall, f1_score = report_dict["accuracy"], report_dict[metrics_average][
            "precision"], report_dict[metrics_average]["recall"], report_dict[metrics_average]["f1-score"]
        self.acc_array.append(cur_acc)
        self.precision_array.append(precision)
        self.recall_array.append(recall)
        self.f1_array.append(f1_score)
        self.best_acc = max(
            cur_acc, self.best_acc) if self.best_acc else cur_acc

        return {
            "cur_acc": cur_acc,
            "best_acc": self.best_acc,
            "acc_array": self.acc_array,
            "precision_array": self.precision_array,
            "recall_array": self.recall_array,
            "f1_array": self.f1_array,
            "classification_report": report_dict,
        }

    def post_message(self, cur_acc, best_acc, output_dir, classification_report,):
        if cur_acc == best_acc:
            write_dict_to_file(classification_report,
                               f"{output_dir}/classification_report.json")


class LossMetrics(ArrayProcessor):
    def message(self, train_loss) -> dict:
        self.append(train_loss)
        return {
            "loss": train_loss,
            "loss_array": self.array
        }


class LRMetrics(ArrayProcessor):
    def message(self, lr) -> dict:
        self.append(lr)
        return {
            "lr": lr,
            "lr_array": self.array
        }


class ROC_AUCMetrics(ArrayProcessor):
    def message(self, y_true, y_score, roc_curve_kwargs, num_classes, output_dir) -> dict:
        y_true = label_binarize(y_true, classes=list(range(num_classes)))
        y_score = np.array(y_score)
        aucs, fprs, tprs = [], [], []
        for i in range(num_classes):
            fpr, tpr, _ = roc_curve(
                y_true[:, i], y_score[:, i], **roc_curve_kwargs)
            aucs.append(auc(fpr, tpr))
            fprs.append(fpr)
            tprs.append(tpr)
        avg_auc = sum(aucs) / len(aucs)
        self.append(avg_auc)

        return {
            "avg_auc": avg_auc,
            "avg_auc_array": self.array,
            "fprs": fprs,
            "tprs": tprs,
            "aucs": aucs,
        }

    def post_message(self, fprs, tprs, aucs, y_true, y_pred, y_score, best_acc, cur_acc, idx_to_class, output_dir):
        if cur_acc == best_acc:
            plt.figure()
            plt.plot([0, 1], [0, 1], 'k--')
            d = {}
            for idx, (fpr, tpr, auc) in enumerate(zip(fprs, tprs, aucs)):
                d[f"{idx_to_class[idx]}_fpr"] = fpr.tolist()
                d[f"{idx_to_class[idx]}_tpr"] = tpr.tolist()
                d[f"{idx_to_class[idx]}_auc"] = auc
                plt.plot(fpr, tpr, label='Class %s (area = %0.2f)' %
                         (idx_to_class[idx], auc))
                plt.legend()
                plt.savefig(f"{output_dir}/roc.png")
            plt.close()
            d["y_true"] = y_true
            d["y_score"] = y_score
            d["y_pred"] = y_pred
            write_dict_to_file(d, f"{output_dir}/roc.json")


class PlotProcessor(Processor):
    def __init__(self):
        super().__init__()
        try:
            from IPython import display
        except ImportError:
            display = None
        self.display = display

    def post_message(self, gdata: dict, total_epochs, dynamic_plot, plot_keys):
        if not dynamic_plot or self.display is None:
            return
        self.display.clear_output(wait=True)
        plt.xlim([0, total_epochs])
        for k, v in gdata.items():
            if k.endswith("_array"):
                k = k.split("_array")[0]
                if k in plot_keys:
                    plt.plot(v, label=k)
        plt.legend()
        plt.pause(0.1)


class ExportProcessor(Processor):
    def __init__(self):
        super().__init__()

    def post_message(self, output_dir, gdata: dict):
        d = {k.split("_array")[0]: v for k,
             v in gdata.items() if k.endswith("_array")}
        df = pd.DataFrame(d)
        df.to_csv(f"{output_dir}/metrics.csv", index_label="epoch")


class ConfusionMatrixProcessor(Processor):
    def post_message(self, y_true, y_pred, output_dir, idx_to_class, cur_acc, best_acc):
        if cur_acc == best_acc:
            classes = [idx_to_class[i] for i in range(len(idx_to_class))]
            cm = confusion_matrix(y_true, y_pred)
            df = pd.DataFrame(cm, index=classes, columns=classes)
            df.to_csv(f"{output_dir}/best_acc_confusion_matrix.csv")


class ModelAutoSaveProcessor(Processor):

    def post_message(self,
                     cur_acc,
                     best_acc,
                     model,
                     output_dir,
                     save_best_model,
                     save_last_model):
        if save_last_model:
            torch.save(model.state_dict(), output_dir / f"last.pth")
        if save_best_model and cur_acc == best_acc:
            sp = pathlib.Path(output_dir, "weights")
            sp.mkdir(parents=True, exist_ok=True)
            torch.save(model.state_dict(), sp / f"best.pth")


class ZipProcessor(Processor):
    def last_message(self, output_dir, auto_zip):
        if auto_zip:
            zip_name = zip_utils.zip_dir(output_dir, output_dir + ".zip")
            print(f"create zip: {zip_name}")


class ProcessorManager:
    def __init__(self, settings: dict):
        self.settings = settings
        self.processors: List[Processor] = [
            ClassificationReportMetrics(),
            LossMetrics(),      # loss, loss_array
            LRMetrics(),        # lr, lr_array
            ROC_AUCMetrics(),
            ConfusionMatrixProcessor(),
            PlotProcessor(),
            ModelAutoSaveProcessor(),
            ExportProcessor(),

            ZipProcessor(),
        ]
        self.keys = self.settings["tqdm_keys"]

    def __process_result(self, data):
        g_data = {}
        for m in self.processors:
            g_data.update(m.distribute_message(data))
        return g_data

    def __post_process_result(self, data: dict):
        for m in self.processors:
            m.distribute_post_message(data)

    def end_current_epoch(self, result):
        coll_dict = {**result.kwargs, **self.settings}
        g_data = self.__process_result(coll_dict)  # settings + result
        coll_dict["gdata"] = g_data
        coll_dict.update(g_data)
        self.__post_process_result(coll_dict)  # settings + result + g_data
        return g_data

    def end_fit(self, extra_config):
        coll_dict = {**self.settings, **extra_config}
        for m in self.processors:
            m.distribute_last_message(coll_dict)
